import pandas as pd
import numpy as np
import xlwt
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import tqdm
import json
import random
import requests
import utils

from openpyxl import load_workbook, Workbook

prompt_2186 = [
    '销售员正在对用户进行360金融产品的营销，希望用户能够进行360金融产品的申请与使用',
    '销售人员正在向用户推销360金融产品，希望用户能够申请并使用',
    '销售人员正在为用户宣传360金融产品，并希望用户能够使用该产品',
    '销售人员正在向用户介绍360金融产品，并鼓励用户进行申请和使用',
    '销售人员正在与用户进行360金融产品的营销，期待用户能够积极参与申请和使用',
    '销售人员正在向用户推介360金融产品，希望用户能够体验和使用该产品',
    '销售人员正在劝说用户尝试使用360金融产品，希望用户能够接受并进行申请',
    '销售人员正在向用户展示360金融产品的优势，并希望用户能够尝试使用',
    '销售人员正在与用户进行360金融产品的推广，鼓励用户进行申请和使用',
    '销售人员正在为用户推销360金融产品，希望用户能够认可并进行申请和使用',
    '销售人员正在向用户介绍360金融产品的特点，并希望用户能够申请并使用该产品'
]
prompt_2199 = [
    '销售员正在对用户进行赠险转短险的营销，用户之前领取过赠险，希望用户能够进行短险的办理',
    '销售人员正在向用户推销短期保险，希望用户能够申请办理',
    '销售人员正在为用户宣传转短期保险的优势，并鼓励用户进行办理',
    '销售人员正在与用户进行赠险转短险的推广，期待用户能够积极参与办理',
    '销售人员正在向用户介绍短期保险的特点，并希望用户能够申请并办理',
    '销售人员正在劝说用户尝试办理短期保险，希望用户能够接受并进行办理',
    '销售人员正在为用户推介短期保险，希望用户能够体验和使用该产品',
    '销售人员正在向用户展示短期保险的优势，并希望用户能够尝试办理',
    '销售人员正在与用户进行短期保险的营销，鼓励用户进行办理',
    '销售人员正在为用户推销转短期保险，希望用户能够认可并进行办理',
    '销售人员正在向用户介绍赠险转短险的优惠政策，并希望用户能够办理短期保险'
]
attitude_dict = {'肯定态度': 1, '不明态度': 0, '否定态度': -1}

model_path = "/mydata/internal_server_model/bloom-560m-nlu-0511"
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

cuda = "cuda:3"
model = model.to(cuda)

def generate_prompt(instruction="baoxian", input=None):
    if instruction == "baoxian":
        instruction = r'下面是段销售员和用户的一次对话。{}。需要对用户的话术进行信息提取，信息包括意图，态度，实体，任何一种信息没有就输出空。'.format(random.choice(prompt_2199))
        # instruction = "根据以下销售员与用户的对话，提取出用户在最后一次表述中的信息，包括意图、态度、实体、标签，每种信息以分号隔开，没有就输出为空，场景为黑牛保险。"
    elif instruction == "jinrong":
        instruction = r"下面是段销售员和用户的一次对话。{}。需要对用户的话术进行信息提取，信息包括意图，态度，实体，任何一种信息没有就输出空。".format(random.choice(prompt_2186))
        # instruction = "根据以下销售员与用户的对话，提取出用户在最后一次表述中的信息，包括意图、态度、实体、标签，每种信息以分号隔开，没有就输出为空，场景为360金融。"
    if input:
        return r"""Instruction:\n{}\n\n### Input:\n{}\n\n### Response:""".format(instruction, input)
    else:
        return r"""Instruction:\n{}\n\n### Response:""".format(instruction)

def inference(instruction, content):
    prompt = generate_prompt(instruction, content)
    # print('content ', content)
    # print('prompt ', prompt)
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(cuda)
    generation_output = model.generate(
        input_ids=input_ids,
        temperature=0.8,
        top_p=0.9,
        top_k=1,
        do_sample=True,
        num_beams=1,
        max_new_tokens=100,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        return_dict_in_generate=True,
        output_scores=True
    )

    s = generation_output.sequences[0]
    output = tokenizer.decode(s)
    res = output.strip().replace(prompt, "").replace("</s>", "")
    # print('response ', res)

    return res

def inference_online(instruction, content):
    url = 'http://39.98.192.246:8679/nlu?session_id=-1&workspace=74743&dialogue={}'

    response = requests.get(url.format(content))
    result = json.loads(response.text)
    attitude = result['original_attitude']

    if result['intention'] and result['intention'][0].get('value', ''):
        intention = result['intention'][0]['value']['standard_query']
        if intention in ['肯定态度', '否定态度', '无态度']:
            intention = ''
    else:
        intention = ''

    return '意图:{};态度:{};实体:{};'.format(intention, attitude, '')



def generate_data(infile, outfile, model_type):


    wb = load_workbook(infile)
    ws = wb[wb.sheetnames[0]]

    wb_w = Workbook()
    ws_w = wb_w.active

    content = r""
    dialogue_num = 0
    for i, row in tqdm.tqdm(enumerate(ws.values)):
        if i == 0:
            ws_w.append(row)
            continue

        if i > 5000:
            break
        '''
        if i < 5000:
            continue
        if i > 13000 :
            break
        '''

        _id = row[0]
        _session = row[1]
        _role = row[2]
        _content = row[3]
        _clean_content = row[4]
        _label = row[5]
        _change_label = row[6]

        if not _id:
            content = r""
            ws_w.append([])
            continue
        else:
            res = ''
            if _role == "销售员" and "用户" in content and model_type == 'model':
                res = inference("baoxian", content)
                # res = inference_online("baoxian", content)
                intent, attitude, entity, _ = res.split(';')
                intent = intent.split(':')[1]
                attitude = attitude.split(':')[1]
                entity = entity.split(':')[1]
                res = '意图:{};态度:{};实体:{};'.format(intent, attitude_dict[attitude], entity)
                # print( content )

            if model_type == 'online':
                res = _label

            label = _change_label if _change_label else _label

            ws_w.append([
                _id, _session, _role, _content, _clean_content, label, res
            ])

            if _role == '销售员':
                content = r"{}:{}\n".format(_role, _clean_content.split('，')[-1])
            else:
                content += r"{}:{}\n".format(_role, _clean_content)

    wb_w.save(outfile)


if __name__ == "__main__":



    # infile = "./dialogue_data/_2186_2023-05-12 14_00_00_2023-05-12 23_30_00_500-13000校对.xlsx"
    # generate_data(infile, "./dialogue_data/2186_model_out.xlsx", 'model')
    # generate_data(infile, "./dialogue_data/2186_online_out.xlsx", 'online')


    infile = './dialogue_data/_2201_2023-05-19 00_00_00_2023-05-19 09_30_00_100000000.xlsx'
    generate_data(infile, "./dialogue_data/2201_model_out.xlsx", 'model')
    generate_data(infile, "./dialogue_data/2201_online_out.xlsx", 'online')






